#!/usr/bin/env python
#-*- coding=utf-8 -*-

#######################################################
# DataSaver负责将QihuiwangPicker抓取的数据保存到文件
# 根据文件的后缀判定保存的文件格式
# *.txt 文本格式
# *.csw 简单电子表格
# *.xls Office电子表格
#######################################################

import os

#如果没有安装xlwt，那么保存为xls文件就不可用
xlwt_enabled = True
try:
    import xlwt
except ImportError: 
    xlwt_enabled = False
    print('xls didn\'t install.')

def save_as_csw(prod_list, filename):
    if len(prod_list) == 0:
        return False

    #分类 商品 价格 联系人 手机 公司 座机 传真 地址 公司网址 源自网页 
    line_fmt = '"%s"\t"%s"\t"%s"\t"%s"\t"%s"\t"%s"\t"%s"\t"%s"\t"%s"\t"%s"\t"%s"\n'
    lines = []
    head_line = line_fmt % ('分类', '商品', '价格', '联系人', '手机号','公司', 
                            '电话', '传真', '公司地址', '公司网址', '源网页')
    lines.append(head_line)
    for item in prod_list:
        info = item['detail']
        if info == None:    #如果信息不全，则跳过
            continue

        prod_line = line_fmt % (item['path'], info['name'], info['price'],\
                                info['contact'], info['cell-phone'], info['company'], \
                                info['tel-phone'], info['fax'], info['address'], info['website'], item['url'])
        lines.append(prod_line)
    
    wfile = open(filename, 'w')
    wfile.writelines(lines)
    wfile.close()
    return True

def save_as_txt(prod_list, filename):
    if len(prod_list) == 0:
        return False

    #分类 商品 价格 联系人 手机 公司 座机 传真 地址 公司网址 源自网页 
    line_fmt = '%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n'
    lines = []
    head_line = line_fmt % ('分类', '商品', '价格', '联系人', '手机号','公司', 
                            '电话', '传真', '公司地址', '公司网址', '源网页')
    lines.append(head_line)
    for item in prod_list:
        info = item['detail']
        if info == None:    #如果信息不全，则跳过
            continue

        prod_line = line_fmt % (item['path'], info['name'], info['price'],\
                                info['contact'], info['cell-phone'], info['company'], \
                                info['tel-phone'], info['fax'], info['address'], info['website'], item['url'])
        lines.append(prod_line)
    
    wfile = open(filename, 'w')
    wfile.writelines(lines)
    wfile.close()
    return True

#保存数据到xls文件中，每个大类放在不同的Worksheet中
def save_as_xls(prod_list, filename):
    if len(prod_list) == 0:
        return False

    workbook = xlwt.Workbook(encoding='utf-8')  #必须注明编码格式，否则会保存失败
    curr_category = ''
    worksheet = None
    row_index = 0
    for prod_item in prod_list:
        path = prod_item['path']
        this_category = path.split('/')[0]
        #如果当前的这个商品种类与上一个商品不同，则要新建worksheet
        if this_category != curr_category:
            worksheet = workbook.add_sheet(this_category)
            curr_category = this_category
            #填写表头
            header_cells = ('分类', '商品', '价格', '联系人', '手机号','公司', 
                     '电话', '传真', '公司地址', '公司网址', '源网页')

            column_index = 0
            for cell in header_cells:
                worksheet.write(0, column_index, header_cells[column_index])
                column_index += 1
            #创建了新了worksheet后，数据从第二行开始往下写
            row_index = 1 

        #将数据填写到worksheet的row_index行
        prod_info = prod_item['detail']
        #如果信息不全，则跳过
        if prod_info == None:
            continue

        prod_cells = (path, prod_info['name'], prod_info['price'], prod_info['contact'],\
                 prod_info['cell-phone'], prod_info['company'], prod_info['tel-phone'],\
                 prod_info['fax'], prod_info['address'], prod_info['website'], prod_item['url'])

        column_index = 0
        for cell in prod_cells:
            worksheet.write(row_index, column_index, prod_cells[column_index])
            column_index += 1

        row_index += 1
        pass
    workbook.save(filename)
    return True
    pass

def get_filename_postfix(filename):
    basename = os.path.basename(filename)
    temp = basename.split('.')
    if len(temp) >= 2:
        return temp[-1]

class DataSaver:
    #后缀名与保存函数映射表
    case_dict = {'csw':save_as_csw,\
                 'txt':save_as_txt}
    if xlwt_enabled:
        case_dict['xls'] = save_as_xls

    #将商品列表数据‘喂’给DataSaver
    def feed(self, data):
        self.product_list = data
        pass

    def save_as(self, filename):
        if self.product_list == None or len(self.product_list) == 0:
            print('警告：记录为空，不保存')
            return

        print('正在保存……')
        while True:
            postfix = get_filename_postfix(filename)
            try:
                if self.case_dict[postfix](self.product_list, filename):
                    print('已保存到：' + filename)
                else:
                    print('保存失败！')
                break
            except KeyError:
                print('警告：不支持 %s 文件格式。' % (postfix))
                print('支持的文件格式：' + ','.join(self.case_dict.keys()))
                try:
                    filename = raw_input('请输入新文件名：')
                except KeyboardInterrupt:
                    print('用户取消保存')
                    break
        pass
    pass

#==============================================================================
def test():
    pl = []
    pl.append({'url':'http://www.xxx.com', 'name':'后悔药', 'path':'医学/药/精神科', \
               'detail':{'name':'后悔药', 'price':'23元', 'detail':'', 'start-date':'', 'end-date':'',\
               'keywords':'', 'contact':'黄博士', 'company':'精神病研究院', 'address':'北京天安门大楼', 'tel-phone':'020-009009',
               'cell-phone':'1343522222', 'fax':'020-399232', 'website':'http://www.shenjinbin.com'}})
    pl.append({'url':'http://www.xyxx.com', 'name':'锣丝钉', 'path':'建材/基本物料/部件', \
               'detail':{'name':'锣丝钉', 'price':'0.2元', 'detail':'', 'start-date':'', 'end-date':'',\
               'keywords':'', 'contact':'刘经理', 'company':'金壁辉煌公司', 'address':'深圳南山', 'tel-phone':'0755-009009',
               'cell-phone':'13409093234', 'fax':'0755-33435232', 'website':'http://www.zhuangxiu.com'}})
    s = DataSaver()
    s.feed(pl)
    s.save_as('test.xxt')
    pass

if __name__ == '__main__':
    test()
